package security.diffie.hellmann;

import java.io.IOException;
import java.math.BigInteger;
import java.util.Random;

import tools.FileUtils;

/**
 * Protocole Diffie-Hellmann.
 * @author Jean-François Augé
 */
public class DiffieHellman {

	private int cores, key_size;
	
	private Random rng;
	private BigInteger p, q, g, A, B;
	private BigInteger a, b, k;
	
	/**
	 * Constructeur.
	 * @param cores
	 * @param key_size
	 */
	public DiffieHellman(int cores, int key_size) {
		this.cores = cores;
		this.key_size = key_size;
		rng = BigInteger_.initRNG();
	}
	
	/**
	 * Générer P = 2Q + 1, avec P et Q premiers.
	 * @return
	 * @throws InterruptedException
	 */
	public BigInteger gen_P() throws InterruptedException {
		do {
			q = Prime.generatePrime(cores, key_size);
			p = q.multiply(BigInteger_.TWO).add(BigInteger.ONE);
		} while(!BigInteger_.fermat(p));
		
		return p;
	}
	
	
	/**
	 * Trouver G, le générateur de P.
	 * @return
	 * @throws InterruptedException
	 */
	public BigInteger gen_G() throws InterruptedException {
		if(!BigInteger_.pow(BigInteger_.TWO, BigInteger_.TWO).equals(BigInteger.ONE.mod(p)))
			g = BigInteger_.TWO;
		else if(!BigInteger_.pow(BigInteger_.THREE, BigInteger_.TWO).equals(BigInteger.ONE.mod(p)))
			g = BigInteger_.THREE;
		else if(!BigInteger_.pow(BigInteger_.FIVE, BigInteger_.TWO).equals(BigInteger.ONE.mod(p)))
			g = BigInteger_.FIVE;
		else if(!BigInteger_.pow(BigInteger_.SEVEN, BigInteger_.TWO).equals(BigInteger.ONE.mod(p)))
			g = BigInteger_.SEVEN;
		
		return g;
	}
	
	/**
	 * Générer la première partie de la clé, A = g^a mod p.
	 * @return
	 * @throws InterruptedException
	 */
	public BigInteger gen_A() throws InterruptedException {
		if(p != null && g != null) {
			a = BigInteger_.generate(rng, key_size);
			return A = g.modPow(a, p);
		}
		else return BigInteger.ZERO;
	}
	
	/**
	 * Générer la deuxième partie de la clé, B = g^b mod p.
	 * @return
	 * @throws InterruptedException
	 */
	public BigInteger gen_B(/*BigInteger p, BigInteger g*/) throws InterruptedException {
		//this.p = p;
		//this.g = g;
		b = BigInteger_.generate(rng, key_size);
		return B = g.modPow(b, p);
	}
	
	/**
	 * Générer la deuxième partie de la clé, B = g^b mod p.
	 * @return
	 * @throws InterruptedException
	 */
	public BigInteger gen_B(BigInteger p, BigInteger g) throws InterruptedException {
		this.p = p;
		this.g = g;
		b = BigInteger_.generate(rng, key_size);
		return B = g.modPow(b, p);
	}
	
	/**
	 * Générer la clé.
	 * @param C
	 * @return
	 */
	public BigInteger compute_key(BigInteger C) {
		if(b != null)
			k = C.modPow(b, p);
		else
			k = C.modPow(a, p);
		
		return k;
	}
	
	/**
	 * Méthode de debug.
	 * @param name
	 */
	public void info(String name) {
		System.out.println(name+" side:");
		if(p != null) System.out.println("p "+getP());
		if(g != null) System.out.println("g "+getG());
		
		if(a != null) System.out.println("a "+a);
		if(A != null) System.out.println("A "+getA());
		if(b != null) System.out.println("b "+b);
		if(B != null) System.out.println("B "+getB());
		
		if(k != null) System.out.println("k "+getK());
	}
	
	public BigInteger getP() { return p; }
	public void setP(BigInteger p) { this.p = p; }

	public BigInteger getQ() { return q; }
	public void setQ(BigInteger q) { this.q = q; }

	public BigInteger getG() { return g; }
	public void setG(BigInteger g) { this.g = g; }

	public BigInteger getA() { return A; }
	public void setA(BigInteger a) { this.A = a; }

	public BigInteger getB() { return B; }
	public void setB(BigInteger b) { this.B = b; }

	public BigInteger getK() { return k; }
	public void setK(BigInteger k) { this.k = k; }

	public int getCores() { return cores; }
	public void setCores(int cores) { this.cores = cores; }

	public int getKey_size() { return key_size; }
	public void setKey_size(int key_size) { this.key_size = key_size; }

	
	// tests
	public static void main(String[] args) {
		int keySize = 1024;
		int cores = 4;
		DiffieHellman alice = new DiffieHellman(cores, keySize);
		DiffieHellman bob = new DiffieHellman(cores, keySize);

		try {
			BigInteger P = alice.gen_P();
			BigInteger G = alice.gen_G();
			BigInteger A = alice.gen_A();
			BigInteger B = bob.gen_B(P, G);

			BigInteger K1 = alice.compute_key(B);
			BigInteger K2 = bob.compute_key(A);

			System.out.println(K1);
			
			CipherModule.writeDHKey("dh.key", K1.toString());

			String test = "une ligne à chiffrer";
			String key = CipherModule.readDHKey("dh.key");
			byte[] test_c = CipherModule.chiffrement_ecb(key, test.getBytes());

			System.out.println(new String(test_c));
			test_c = CipherModule.chiffrement_ecb(key, test_c);
			System.out.println(new String(test_c));
			
			System.out.println("Test fichier !");
			FileUtils.copyFileBufferedEncryptedDH("C:\\Documents and Settings\\Guillaume\\Mes documents\\Visual Studio 2010\\Projects\\test_compatibilite_crypto\\test_compatibilite_crypto\\key\\a51.key",
					"C:\\Documents and Settings\\Guillaume\\Mes documents\\Visual Studio 2010\\Projects\\test_compatibilite_crypto\\test_compatibilite_crypto\\key\\a51.key_c", 4, K1.toString());
			FileUtils.copyFileBufferedEncryptedDH("C:\\Documents and Settings\\Guillaume\\Mes documents\\Visual Studio 2010\\Projects\\test_compatibilite_crypto\\test_compatibilite_crypto\\key\\a51.key_c",
					"C:\\Documents and Settings\\Guillaume\\Mes documents\\Visual Studio 2010\\Projects\\test_compatibilite_crypto\\test_compatibilite_crypto\\key\\a51.key_", 4, K2.toString());
		} catch (InterruptedException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		/*String test = "/etc/kraken/key/a51.key_c";
		System.out.println(test+"\n"+test.substring(0,test.length()-2));*/

	}

}
